import os
import sys
import re
import datetime

import numpy as np
import torch
# from torch.optim.lr_scheduler import _LRScheduler
import torchvision
import torchvision.transforms as transforms
from torch.utils.data import DataLoader

"""
Reference: 
https://github.com/weiaicunzai/pytorch-cifar100

"""


def get_work(args):
    if args.net == 'vgg16':
        from VGG.vgg import vgg16_bn
        net = vgg16_bn()
    elif args.net == 'vgg13':
        from VGG.vgg import vgg13_bn
        net = vgg13_bn()
    elif args.net == 'vgg11':
        from VGG.vgg import vgg11_bn
        net = vgg11_bn()
    elif args.net == 'vgg19':
        from VGG.vgg import vgg19_bn
        net = vgg19_bn()
    else:
        raise ValueError('Wrong Network Name')

    return net


def get_training_dataloader(mean, std, batch_size=16, num_workers=2, shuffle=True):
    """
    - Inputs:
        mean: mean of cifar100 training dataset
        std: std of cifar100 training dataset
        batch_size: dataloader batch size
        num_workers: dataloader #workers
        shuffle: whether to shuffle
    - Returns
        torch dataloader
    """

    transform_train = transforms.Compose([
        transforms.RandomCrop(32, padding=4),
        transforms.RandomHorizontalFlip(),
        transforms.RandomRotation(15),
        transforms.ToTensor(),
        transforms.Normalize(mean, std)
    ])

    cifar100_training = torchvision.datasets.CIFAR100(root='./data', train=True, download=True,
                                                      transforms=transform_train)
    cifar100_training_loader = DataLoader(
        cifar100_training, shuffle=shuffle, num_workers=num_workers, batch_size=batch_size
    )

    return cifar100_training_loader


def get_test_dataloader(mean, std, batch_size=16, num_workers=2, shuffle=True):
    transform_test = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize(mean, std)
    ])

    cifar100_test = torchvision.datasets.CIFAR100(root='./data', train=False, download=True, transform=transform_test)
    cifar100_test_dataloader = DataLoader(
        cifar100_test, shuffle=shuffle, num_workers=num_workers, batch_size=batch_size
    )

    return cifar100_test_dataloader


def compute_mean_std(cifar100_dataset):
    data_r = np.dstack([cifar100_dataset[i][1][..., 0]] for i in range(len(cifar100_dataset)))
    data_g = np.dstack([cifar100_dataset[i][1][..., 1]] for i in range(len(cifar100_dataset)))
    data_b = np.dstack([cifar100_dataset[i][1][..., 2]] for i in range(len(cifar100_dataset)))

    mean = np.mean(data_r), np.mean(data_g), np.mean(data_b)
    std = np.std(data_r), np.std(data_g), np.std(data_b)

    return mean, std


def most_recent_weights(weights_folder):
    weight_files = os.listdir(weights_folder)
    if len(weight_files) == 0:
        return str()
    regex_str = r'([A-Za-z0-9]+)-([0-9]+)-(regular|best)'

    weight_files = sorted(weight_files, key=lambda w: int(re.search(regex_str, w).groups()[1]))

    return weight_files[-1]


if __name__ == '__main__':
    a = np.array((1, 2, 3))
    b = np.array((2, 3, 4))
    ret = np.dstack((a, b))
    res = np.vstack((a, b))
    d = np.random.random((3, 4, 5))
    e = d[..., 1]
    f = d[:, :, 1]
    m = e == f
    pass
    pass
